import numpy as np  
import matplotlib.pyplot as plt  
from filterpy.kalman import ExtendedKalmanFilter as EKF  
from scipy.spatial.transform import Rotation as R  
  
import pandas as pd  
  

# 生成螺旋上升运动的三维坐标数据  
def generate_spiral():  
    """Generate coordinates for a spiral motion."""  
    coordinates = []  
    t = np.linspace(0, 10, int(10 * 30), endpoint=False)  # 0到10秒，30Hz  
    for ti in t:  
        # 螺旋线的参数方程  
        x = amplitude * np.cos(ti)  
        y = amplitude * np.sin(ti)  
        z = amplitude * ti  # 时间作为z轴（上升） 
        coordinates.append([x, y, z])  
    return np.array(coordinates)  
  
# 添加高斯噪声  
def add_gaussian_noise(data, noise_sigma=10):  
    """Add Gaussian noise to the data."""  
    noisy_data = data + np.random.normal(0, noise_sigma, data.shape)  
    return noisy_data  
  
# 扩展卡尔曼滤波器模型定义  
def f(x, dt):  
    """State transition function for nearly constant velocity."""  
    return np.array([x[0] + x[3]*dt, x[1] + x[4]*dt, x[2] + x[5]*dt, x[3], x[4], x[5]])   
  
def h(x):  
    """Measurement function."""
    return np.array([x[0], x[1], x[2]])  
  
def F(x, dt):  
    """State transition jacobian."""  
    return  np.array([1, 0, 0, dt, 0, 0],  
        [0, 1, 0, 0, dt, 0],  
        [0, 0, 1, 0, 0, dt],  
        [0, 0, 0, 1, 0, 0],  
        [0, 0, 0, 0, 1, 0],  
        [0, 0, 0, 0, 0, 1] )
  
def H(x):  
    """Measurement jacobian."""  
    return np.array([[1, 0, 0, 0, 0, 0],  
                     [0, 1, 0, 0, 0, 0],  
                     [0, 0, 1, 0, 0, 0]])  

# 读取Excel文件  
file_path = 'vision_pos.xlsx'  # 替换为你的Excel文件路径  
df = pd.read_excel(file_path, sheet_name='Sheet1')  
  
# 初始化一个空列表来存储坐标  
coordinates = []  
  
# 遍历DataFrame的每一行，将数据添加到coordinates列表中  
for index, row in df.iterrows():  
    x = row.iloc[0]  # 假设第一项是x坐标  
    y = row.iloc[1]  # 假设第二项是y坐标  
    z = row.iloc[2]  # 假设第三项是z坐标  
    coordinates.append([x, y, z])  
  
# 将列表转换为numpy数组  
noisy_data=np.array(coordinates[:939])

# 设置时间、频率和振幅参数  
t_total = 94  # 总时间（秒）  
frequency = 10  # 频率（Hz）

amplitude = 100  # 振幅（厘米），为了看到明显的螺旋线，这里选择较大的振幅值  
dt = 1 / frequency  # 时间步长（秒）  
num_points = int(t_total / dt)  # 总点数  
times = np.arange(0, t_total, dt)  # 时间数组  
  

# -5.74
# 1.23
# 0.46
qconstant = 6
rconstant = 125
# sequence
# 设置扩展卡尔曼滤波器参数并初始化滤波器状态向量和协方差矩阵  
ekf = EKF(dim_x=6, dim_z=3)  # 状态向量为[x, y, z]，测量向量为[x, y, z]  
ekf.x = np.array([noisy_data[0,0], noisy_data[0,1], noisy_data[0,2], 0, 0, 0])  # 初始状态（位置和速度）设置为第一个测量值，速度为0  
ekf.P = np.eye(6) * 0**2  # 初始状态协方差矩阵，对角线上的值表示对应状态的不确定性程度（这里假设所有状态的不确定性相同）  

ekf.R = np.eye(3) * (4*rconstant)**2   # 测量噪声协方差矩阵，设置为噪声的标准差的平方（转换为厘米） 
ekf.Q = np.eye(6) * (qconstant**2)  # 过程噪声协方差矩阵，这里假设过程噪声较小，因此设置为较小的值（可以根据实际情况调整）  
ekf.F = np.array([[1, 0, 0, 0, 0, 0],  
                  [0, 1, 0, 0, 0, 0],  
                  [0, 0, 1, 0, 0, 0],  
                  [0, 0, 0, 1, 0, 0],  
                  [0, 0, 0, 0, 1, 0],  
                  [0, 0, 0, 0, 0, 1]]) # 状态转移函数的雅可比矩阵（使用lambda函数传递dt）  
ekf.H = H  # 测量函数的雅可比矩阵（不需要状态向量，因此直接传递函数）  
ekf.f = f  # 状态转移函数（使用lambda函数传递dt）  
ekf.h = h  # 测量函数（不需要状态向量，因此直接传递函数）  
#ekf.R[:] = (10*100)**2  # 更新测量噪声协方差矩阵的所有元素为相同值（因为我们认为所有方向的测量噪声都相同）

filtered_data_nom = []  # 用于存储滤波后的数据  

delt_nom = []

for tt in noisy_data:  # 对每个测量值进行滤波处理  
    ekf.update(tt,HJacobian=ekf.H,Hx=ekf.h)
      # 使用当前测量值更新滤波器状态估计和协方差矩阵  
    ekf.predict()  # 使用状态转移函数预测下一个时间步的状态估计和协方差矩阵（尽管这里我们实际上是在处理当前时间步的下一个测量值）
    filtered_data_nom.append(ekf.x)  # 存储滤波后的位置估计（忽略速度估计）  
    delt_nom.append(np.linalg.norm(ekf.x[:3] - tt))
filtered_data_nom = np.array(filtered_data_nom)  # 将列表转换为NumPy数组以便于绘图处理  
delt_nom = np.array(delt_nom)

fig = plt.figure()  

# Cam
ax1 = fig.add_subplot(111, projection='3d')  
ax1.plot(filtered_data_nom[:, 0], filtered_data_nom[:, 1], filtered_data_nom[:, 2], 'grey', label='Original')  # 原图为蓝色  
ax1.set_xlabel("X Position (cm)")  
ax1.set_ylabel("Y Position (cm)")  
ax1.set_zlabel("Z Position (cm)")  
ax1.set_title("Trajectory after ekf")  
plt.legend()  

# -6 1 0.5
my_list = [[-6, 1, 0.5,0,0,0] for _ in range(938)]  
ll=np.array(my_list)
delt_old=np.array([noisy_data[0,0], noisy_data[0,1], noisy_data[0,2], 0, 0, 0])-ll
sum_old=0
for i in delt_old:
    sum_old=sum_old+i[0]*i[0]+i[1]*i[1]+i[2]*i[2]

delt=filtered_data_nom-ll
fx=np.nanvar(delt, axis=1)
sum_=0

sum_x=0
sum_y=0
sum_z=0
for i in noisy_data:
    sum_x=sum_x+i[0]
    sum_y=sum_y+i[1]
    sum_z=sum_z+i[2]
mean_x=sum_x/938
mean_y=sum_y/938
mean_z=sum_z/938
print("old mean")
my_list = [[mean_x, mean_y, mean_z,0,0,0] for _ in range(938)] 
sum_=0 
for i in noisy_data:
    sum_=sum_+(i[0]-mean_x)*(i[0]-mean_x)+(i[1]-mean_y)*(i[1]-mean_x)+(i[2]-mean_z)*(i[2]-mean_z)
print("old sum")
print(sum_/937)
sum_x_new=0
sum_y_new=0
sum_z_new=0
for i in filtered_data_nom:
    sum_x_new=sum_x_new+i[0]
    sum_y_new=sum_y_new+i[1]
    sum_z_new=sum_z_new+i[2]
mean_x_new=sum_x_new/938
mean_y_new=sum_y_new/938
mean_z_new=sum_z_new/938
print("new mean")
my_list_new = [[mean_x_new, mean_y_new, mean_z_new,0,0,0] for _ in range(938)] 
sum_new=0 
for i in filtered_data_nom:
    sum_new=sum_new+(i[0]-mean_x_new)*(i[0]-mean_x_new)+(i[1]-mean_y_new)*(i[1]-mean_x_new)+(i[2]-mean_z_new)*(i[2]-mean_z_new)
print("new sum")
print(sum_new/937)
print("new delt")
print(mean_x_new+6)
print(mean_y_new-1)
print(mean_z_new-0.5)
# -5.74
# 1.23
# 0.46
plt.show()